#Copyright (C) Nial Peters 2013
#
#This file is part of pydoas.
#
#pydoas is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#pydoas is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with pydoas.  If not, see <http://www.gnu.org/licenses/>.
"""
Module contains reference spectra for DOAS fitting.
"""
import numpy
import os.path
import math
import scipy.interpolate
from scipy.ndimage.filters import gaussian_filter1d
import scipy.signal
import doas
import doas.spec


def __load_conv_solar(fwhm):
    #load the raw reference data
    arr = _load_ref_spec('solar')
    
    conv_counts = convolve_spec(arr[0,:], arr[1,:], fwhm)
    
    return arr[0,:],conv_counts 
     
        
def __load_conv_o3(fwhm):
    #load the raw reference data
    arr = _load_ref_spec('o3')
    
    conv_counts = convolve_spec(arr[0,:], arr[1,:], fwhm)
    
    return arr[0,:],conv_counts


def __load_conv_so2(fwhm):
    #load the raw reference data
    arr = _load_ref_spec('so2')
    
    conv_counts = convolve_spec(arr[0,:], arr[1,:], fwhm, deconv_fwhm=0.21)
    
    return arr[0,:],conv_counts


def __load_conv_ring(fwhm):
    #load the raw reference data
    arr = _load_ref_spec('ring')
    
    conv_raman = convolve_spec(arr[0,:], arr[1,:], fwhm, deconv_fwhm=0.21)
    conv_solar = convolve_spec(arr[0,:], arr[2,:], fwhm, deconv_fwhm=0.21)
    
    ring = numpy.zeros_like(conv_solar)
    mask = numpy.nonzero(conv_solar)
    ring[mask] = conv_raman[mask]/conv_solar[mask]
    
    return arr[0,:],ring



ref_to_filename_map = {'solar':'solar_chance_2010.npy',
                       'ring':'ring_g021.npy',
                       'so2':'so2_scia_e293Kair.npy',
                       'o3':'o3_246K_hp_air_bremen.npy'}

ref_to_loader_func_map = {'solar':__load_conv_solar,
                          'ring':__load_conv_ring,
                          'so2':__load_conv_so2,
                          'o3':__load_conv_o3}



def _load_ref_spec(name):
    """
    Loads the named reference spectrum from file and returns it as a 
    numpy array of dimensions (2,n). The first row is the wavelengths, 
    the second is the counts.
    """
    if not name in ref_to_filename_map.keys():
        raise ValueError("Unknown reference spectrum \'%s\' requested. "
                         "Expecting one of: "
                         "%s"%(name, ', '.join(ref_to_filename_map.keys())))
        
    full_path = os.path.join(doas.get_doas_sys_dir(), 'references', 
                             ref_to_filename_map[name])
    
    return numpy.load(full_path)
    


def convolve_spec(wavelengths, counts, conv_fwhm, deconv_fwhm=0.0):
    """
    Deconvolves the spectrum with a Gaussian of FWHM given by deconv_fwhm and
    then convolves it with a Gaussian of FWHM given by conv_fwhm. wavelengths
    and counts should both be numpy arrays.
    """
  
    #calculate the wavelength resolution of the reference spectrum
    diff = wavelengths[1:] - wavelengths[:-1]
    
    #sanity check - we expect the wavelength resolution of the reference to be
    #almost constant (since we are not going to interpolate it onto a regular
    #grid). Make sure this is actually the case - we'll allow a 0.01% difference
    if diff.max() - diff.min() > (diff.max()/10000.0):
        raise ValueError("The reference spectrum has an irregular wavelength "
                         "sampling. This is probably a bug and you should report it!")
    
    resolution = numpy.mean(diff)
    
    #find the std-dev in pixels for the instrument function
    conv_fwhm_pix = conv_fwhm / resolution
    
    #standard relationship between sigma and fwhm
    conv_sigma = conv_fwhm_pix / (2.0 * math.sqrt(2.0 * math.log(2.0)))
    
    #take account of any deconvolution that needs to be done
    if deconv_fwhm > 0:
        deconv_fwhm_pix = deconv_fwhm / resolution
        deconv_sigma = deconv_fwhm_pix / (2.0 * math.sqrt(2.0 * math.log(2.0)))
        
        conv_sigma = math.sqrt(conv_sigma**2 - deconv_sigma**2)
        
    filter = scipy.signal.gaussian(2001, conv_sigma)
    convolved_counts = scipy.signal.fftconvolve(counts, filter,'same')/numpy.sum(filter)
    
    return convolved_counts
    


def get_reference_spec(name, wavelengths, fwhm):
    """
    Returns the named reference spectrum, convolved with a gaussian slit 
    function with a Full-Width at Half Maximum (FWHM) of fmhw and interpolated 
    onto the specified wavelengths using a Cubic Spline.
    
    name may be one of: 'solar', 'ring', 'so2', or 'o3'
    wavelengths must be a numpy array of desired wavelengths.   
    """      
    
    assert fwhm >= 0.21, "pydoas cannot provide references for very high resolution spectrometers."
    
    ref_wavelengths, convolved_counts = ref_to_loader_func_map[name](fwhm)
    
    #sanity check - make sure that the reference actually has some data within
    #the desired wavelength range
    if (wavelengths[0] >= ref_wavelengths[-1] or 
        wavelengths[-1]<ref_wavelengths[0]):
        raise ValueError("Specified wavelength range is outside of the range "
                         "for this reference. The %s reference spans a "
                         "wavelength range of "
                         "%f-%f."%(name, ref_wavelengths[0],ref_wavelengths[-1]))
    
    #if the requested wavelengths extend beyond that of the reference, zero pad
    #the reference to stop the spline from diverging at the ends
    if wavelengths[0] < ref_wavelengths[0]:
        resolution = numpy.mean(wavelengths[1:] - wavelengths[:-1])
        #pre-pad the data
        extra_points_wavelengths = numpy.arange(wavelengths[0], 
                                                ref_wavelengths[0], resolution)
        extra_points_counts = numpy.zeros_like(extra_points_wavelengths)
        ref_wavelengths = numpy.array(numpy.concatenate((extra_points_wavelengths,
                                                         ref_wavelengths)))
        convolved_counts = numpy.array(numpy.concatenate((extra_points_counts, 
                                                          convolved_counts)))
    
    if wavelengths[-1] > ref_wavelengths[-1]:
        resolution = numpy.mean(wavelengths[1:] - wavelengths[:-1])
        #post-pad the data
        extra_points_wavelengths = numpy.arange(ref_wavelengths[-1] + resolution, 
                                                wavelengths[-1] + resolution, 
                                                resolution)
        extra_points_counts = numpy.zeros_like(extra_points_wavelengths)
        ref_wavelengths = numpy.array(numpy.concatenate((ref_wavelengths, 
                                                         extra_points_wavelengths)))
        convolved_counts = numpy.array(numpy.concatenate((convolved_counts, 
                                                          extra_points_counts)))
    
    #use cubic spline interpolation to map the reference onto the wavelengths
    #requested
    interpolator = scipy.interpolate.UnivariateSpline(ref_wavelengths, 
                                                      convolved_counts, s=0)
    interp_counts = interpolator(wavelengths)
    
    return ReferenceSpectrum(wavelengths, interp_counts)
    
    
    
class ReferenceSpectrum:
    def __init__(self, wavelengths, counts):
        self.wavelengths = wavelengths
        self.counts = counts
        
    def save(self, filename):
        """
        Save the reference spectrum as a two column, tab-separated text file
        with the first column containing the wavelengths, and the second 
        containing the reference counts.
        """
        with open(filename, 'w') as ofp:
            for i in range(len(self.wavelengths)):
                ofp.write("%f\t%e\n"%(self.wavelengths[i],self.counts[i]))

